from src.tools.dms_operator import dmsoperator
from src.tools.mml_to_latex import mathml2latex
from src.edit_file.to_document import to_document_bpr
from src.edit_file.db import faiss_embedding1, faiss_embedding2
from src.edit_file.prompt import Prompt
from src.tools.chain import gen_chain
from src.tools.split_string import split_sub_string
from src.tools.extract_jsonstring import re_json_string
from src.tools.information_exchange import pub_gen
import json

class EditBPRFile:

    def __init__(self):
        self.prompt = Prompt()

    # 解析表格
    def parse_table(self, table: dict):
        table_dict = {
            "index": table["index"],
            "sectionIndex": table["sectionIndex"],
            "bodyChildIndex": table["bodyChildIndex"],
            "tableCellInfos": []
        }
        for cell in table["tableCellInfos"]:
            cell_dict = {
                "rowIndex": cell["rowIndex"],
                "columnIndex": cell["columnIndex"],
                "paragraphInfos": [],
            }

            for para in cell["paragraphInfos"]:
                para_text = ""
                for textinfo in para["textInfos"]:
                    if textinfo["subSuperscript"] == 2:
                        para_text += f'<sub>{textinfo["text"]}</sub>'
                    else:
                        para_text += textinfo["text"]

                mathinfos = []
                for mathinfo in para["officeMathInfos"]:
                    math_dict = {
                        "text": mathml2latex(mathinfo["mathMlCode"])
                    }
                    mathinfos.append(math_dict)

                para_dict = {
                    "index": para["index"],
                    "sectionIndex": para["sectionIndex"],
                    "bodyChildIndex": para["bodyChildIndex"],
                    "text": para_text,
                    "latexMathInfos": mathinfos
                }
                cell_dict["paragraphInfos"].append(para_dict)
            table_dict["tableCellInfos"].append(cell_dict)
        return table_dict

    def table_diff(self, rowindex, columnindex, before_tableinfo, after_tableinfo):
        cell_text = ""
        res_before_cell = None
        res_after_cell = None
        for cell in before_tableinfo["tableCellInfos"]:
            if cell["rowIndex"] == rowindex and cell["columnIndex"] == columnindex:
                cell_text = "\n".join([para["text"] for para in cell["paragraphInfos"]])
                res_before_cell = cell
        for after_cell in after_tableinfo["tableCellInfos"]:
            if after_cell["rowIndex"] == rowindex and after_cell["columnIndex"] == columnindex:
                res_after_cell = after_cell
        return cell_text, res_before_cell, res_after_cell

    def edit_cell(self, revise_BPRtext_dict, table, file_pubinfo):
        table_dict = {
            "index": table["index"],
            "sectionIndex": table["sectionIndex"],
            "bodyChildIndex": table["bodyChildIndex"],
            "operationType": 0,
            "tableInfo": {
                "cellEditInfos": []
            }
        }
        cell_dict = {
            "rowIndex": revise_BPRtext_dict["rowIndex"],
            "columnIndex": revise_BPRtext_dict["columnIndex"],
            "cellChildInfos": []
        }
        for para in revise_BPRtext_dict["paragraphInfos"]:
            edit_para_dict = {
                "index": para["index"],
                "sectionIndex": para["sectionIndex"],
                "bodyChildIndex": para["bodyChildIndex"],
                "operationType": 1,
                "paragraphInfo": {
                    "paragraphChildInfos": [],
                    "referenceParagraphType": 2
                }
            }
            # 插入纯文本
            textinfo_list = split_sub_string(para["text"])
            textinfos = []
            for textinfo_index, textinfo in enumerate(textinfo_list):
                textinfo_dict = {
                    "paragraphChildIndex": textinfo_index,
                    "paraChildType": 0,
                    "content": textinfo.get("content"),
                    "subSuperScript": textinfo.get("subSuperScript")
                }
                textinfos.append(textinfo_dict)
            if textinfos:
                edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += textinfos
            # 插入数学式
            if para["latexMathInfos"]:
                latex_index = len(textinfos)
                mathinfos = []
                for latex_math_info in para["latexMathInfos"]:
                    mathinfo_dict = {
                        "paragraphChildIndex": latex_index,
                        "paraChildType": 2,
                        "content": latex_math_info["text"]
                    }
                    latex_index += 1
                    mathinfos.append(mathinfo_dict)
                edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += mathinfos
            cell_dict["cellChildInfos"].append(edit_para_dict)
            # 删除段落，要判断是否在原文档段落index范围内
            if para["index"] < len(table["tableCellInfos"][0]["paragraphInfos"]):
                dele_para_dict = {
                    "index": para["index"],
                    "sectionIndex": para["sectionIndex"],
                    "bodyChildIndex": para["bodyChildIndex"],
                    "operationType": 2
                }
                cell_dict["cellChildInfos"].append(dele_para_dict)
        # 如果生成的内容的index少于原内容的index，则删除剩余的段落
        if len(revise_BPRtext_dict["paragraphInfos"]) < len(table["tableCellInfos"][0]["paragraphInfos"]):
            for cell_para in table["tableCellInfos"][0]["paragraphInfos"]:
                if cell_para["index"] >= len(revise_BPRtext_dict["paragraphInfos"]):
                    dele_para_dict = {
                        "index": cell_para["index"],
                        "sectionIndex": cell_para["sectionIndex"],
                        "bodyChildIndex": cell_para["bodyChildIndex"],
                        "operationType": 2
                    }
                    cell_dict["cellChildInfos"].append(dele_para_dict)

        table_dict["tableInfo"]["cellEditInfos"].append(cell_dict)

        file_pubinfo["editInfos"].append(table_dict)



    # 找到流式输出内容
    def find_key_values(self, target_key, data):
        found_values = []
        if isinstance(data, dict):
            for key, value in data.items():
                if key == target_key and value:
                    found_values.append(value)
                found_values.extend(self.find_key_values(target_key, value))
        elif isinstance(data, list):
            for item in data:
                found_values.extend(self.find_key_values(target_key, item))
        return found_values

    def edit_file(self, revise_file, template_file_path, task_id):
        revise_file_path = revise_file.get("realPath")
        revise_jsonfile = dmsoperator.get_jsonfile(revise_file_path)
        docus_para, para_ids, docus_cell, cell_ids = to_document_bpr(revise_jsonfile, revise_file)
        print(para_ids)
        print(cell_ids)
        file_pubinfo = {
            "filePath": revise_jsonfile["filePath"],
            "isHighlight": False,
            "editInfos": []
        }
        # 修改段落
        if docus_para:
            faiss_embedding1.add_documents(docus_para, ids=para_ids)
        retr_para = faiss_embedding1.as_retriever(search_kwargs={"k": 1})
        template_diffjson = dmsoperator.get_diff_jsonfile(template_file_path)
        # print("-" * 200)
        # file_pubinfo_json = json.dumps(template_diffjson, indent=4, ensure_ascii=False)
        # print(file_pubinfo_json)
        # # 保存修改的json文件
        # edit_SOPjsonfile_path = f"段落修改1.json"
        # with open(edit_SOPjsonfile_path, "w", encoding="utf-8") as f:
        #     f.write(file_pubinfo_json)
        # print("-" * 200)
        for diff_chunk in template_diffjson["data"]:
            for diff_content in diff_chunk["detailDiffDtoList"]:
                # 编辑类型
                if diff_content["operationType"] == 0:
                    revise_retrpara_list = retr_para.invoke(diff_content["beforeContent"])
                    revise_retrpara = revise_retrpara_list[0]
                    faiss_embedding1.delete([revise_retrpara.metadata["id"]])
                    revise_query = self.prompt.edit_para_prompt(diff_content["beforeContent"], diff_content["afterContent"], revise_retrpara.page_content)
                    revise_ans = gen_chain(revise_query)
                    pub_gen(task_id, revise_ans, revise_file.get("docId"), revise_file.get("taskDocId"))
                    edit_para_dict = {
                        "index": revise_retrpara.metadata["index"],
                        "sectionIndex": revise_retrpara.metadata["sectionIndex"],
                        "bodyChildIndex": revise_retrpara.metadata["bodyChildIndex"],
                        "operationType": 1,
                        "paragraphInfo": {
                            "paragraphChildInfos": [],
                            "referenceParagraphType": 2
                        }
                    }
                    # 插入纯文本
                    textinfo_list = split_sub_string(revise_ans)
                    textinfos = []
                    for textinfo_index, textinfo in enumerate(textinfo_list):
                        textinfo_dict = {
                            "paragraphChildIndex": textinfo_index,
                            "paraChildType": 0,
                            "content": textinfo.get("content"),
                            "subSuperScript": textinfo.get("subSuperScript")
                        }
                        textinfos.append(textinfo_dict)
                    if textinfos:
                        edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += textinfos
                    file_pubinfo["editInfos"].append(edit_para_dict)
                    # 删除段落
                    dele_edit_dict = {
                        "index": revise_retrpara.metadata["index"],
                        "sectionIndex": revise_retrpara.metadata["sectionIndex"],
                        "bodyChildIndex": revise_retrpara.metadata["bodyChildIndex"],
                        "operationType": 2
                    }
                    file_pubinfo["editInfos"].append(dele_edit_dict)
                # 插入类型
                if diff_content["operationType"] == 1:
                    revise_retrpara_list = retr_para.invoke(diff_chunk["previousParagraph"])
                    revise_retrpara = revise_retrpara_list[0]
                    edit_para_dict = {
                        "index": revise_retrpara.metadata["index"],
                        "sectionIndex": revise_retrpara.metadata["sectionIndex"],
                        "bodyChildIndex": revise_retrpara.metadata["bodyChildIndex"] + 1,
                        "operationType": 1,
                        "paragraphInfo": {
                            "paragraphChildInfos": [],
                        }
                    }
                    pub_gen(task_id, diff_content["afterContent"], revise_file.get("docId"), revise_file.get("taskDocId"))
                    textinfo_list = split_sub_string(diff_content["afterContent"])
                    textinfos = []
                    for textinfo_index, textinfo in enumerate(textinfo_list):
                        textinfo_dict = {
                            "paragraphChildIndex": textinfo_index,
                            "paraChildType": 0,
                            "content": textinfo.get("content"),
                            "subSuperScript": textinfo.get("subSuperScript")
                        }
                        textinfos.append(textinfo_dict)
                    if textinfos:
                        edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += textinfos
                    file_pubinfo["editInfos"].append(edit_para_dict)

                # 删除类型
                if diff_content["operationType"] == 2:
                    revise_retrpara_list = retr_para.invoke(diff_content["beforeContent"])
                    revise_retrpara = revise_retrpara_list[0]
                    faiss_embedding1.delete([revise_retrpara.metadata["id"]])
                    dele_edit_dict = {
                        "index": revise_retrpara.metadata["index"],
                        "sectionIndex": revise_retrpara.metadata["sectionIndex"],
                        "bodyChildIndex": revise_retrpara.metadata["bodyChildIndex"],
                        "operationType": 2
                    }
                    file_pubinfo["editInfos"].append(dele_edit_dict)

        # 修改表格
        faiss_embedding2.add_documents(docus_cell, ids=cell_ids)
        retr_cell = faiss_embedding2.as_retriever(search_kwargs={"k": 1})
        template_diffjson = dmsoperator.get_diff_jsonfile_table(template_file_path)
        # print("-"*200)
        # file_pubinfo_json = json.dumps(template_diffjson, indent=4, ensure_ascii=False)
        # print(file_pubinfo_json)
        # # 保存修改的json文件
        # edit_SOPjsonfile_path = f"表格修改1.json"
        # with open(edit_SOPjsonfile_path, "w", encoding="utf-8") as f:
        #     f.write(file_pubinfo_json)
        # print("-"*200)
        for diff_chunk in template_diffjson["data"]:
            if diff_chunk["tableDiffInfo"] is None:
                continue
            table_diffinfo = diff_chunk["tableDiffInfo"]
            if table_diffinfo["operationType"] == 1 or table_diffinfo["operationType"] == 2:
                continue
            # 编辑单元格
            before_tableinfo = self.parse_table(table_diffinfo["beforeTableInfo"])
            after_tableinfo = self.parse_table(table_diffinfo["afterTableInfo"])
            for tablecelldiffinfo in table_diffinfo["tableCellDiffInfos"]:
                retr_query, before_cellinfo, after_cellinfo = self.table_diff(tablecelldiffinfo["rowIndex"], tablecelldiffinfo["columnIndex"], before_tableinfo, after_tableinfo)
                if retr_query == "" or before_cellinfo is None or after_cellinfo is None:
                    continue
                revise_retrtable_list = retr_cell.invoke(retr_query)
                revise_retrcell = revise_retrtable_list[0]
                print(revise_retrcell.metadata["id"])
                faiss_embedding2.delete([revise_retrcell.metadata["id"]])
                print(f"检索内容：{revise_retrcell}")
                beforecell_jsontext = json.dumps(before_cellinfo, ensure_ascii=False)
                aftercell_jsontext = json.dumps(after_cellinfo, ensure_ascii=False)
                revise_retrcell_jsontext = json.dumps(revise_retrcell.metadata["tableCellInfos"][0], ensure_ascii=False)
                revise_query = self.prompt.edit_cell_prompt(beforecell_jsontext, aftercell_jsontext, revise_retrcell_jsontext)
                print(f"检索的问题：{revise_query}")
                revise_ans = gen_chain(revise_query)
                print(f"生成内容：{revise_ans}")
                revise_ans_textjson = re_json_string(revise_ans)
                print(f"生成json内容：{revise_ans_textjson}")
                revise_ans_dict = json.loads(revise_ans_textjson)
                text_key = self.find_key_values("text", revise_ans_dict)
                for msg in text_key:
                    pub_gen(task_id, msg, revise_file.get("docId"), revise_file.get("taskDocId"))
                self.edit_cell(revise_ans_dict, revise_retrcell.metadata, file_pubinfo)

            # # 插入表格
            # if table_diffinfo["operationType"] == 1:
            #     if table_diffinfo["previousParagraph"] is None:
            #         continue
            #     revise_retrpara_list = retr_para.invoke(table_diffinfo["previousParagraph"])
            #     revise_retrpara = revise_retrpara_list[0]
            #     aftertable_dict, aftertable_text = self.parse_table(table_diffinfo["afterTableInfo"])
            #     text_key = self.find_key_values("text", aftertable_dict)
            #     for msg in text_key:
            #         pub_gen(task_id, msg, revise_file.get("docId"), revise_file.get("taskDocId"))
            #     self.insert_table(aftertable_dict, revise_retrpara.metadata, file_pubinfo)
            # # 删除表格
            # if table_diffinfo["operationType"] == 2:
            #     beforetable_dict, beforetable_text = self.parse_table(table_diffinfo["beforeTableInfo"])
            #     revise_retrtable_list = retr_table.invoke(beforetable_text)
            #     revise_retrtable = revise_retrtable_list[0]
            #     dele_table = {
            #         "index": revise_retrtable.metadata["index"],
            #         "sectionIndex": revise_retrtable.metadata["sectionIndex"],
            #         "bodyChildIndex": revise_retrtable.metadata["bodyChildIndex"],
            #         "operationType": 2,
            #     }
            #     file_pubinfo["editInfos"].append(dele_table)

        edit_res = dmsoperator.post_edit_file_v2(file_pubinfo)
        res = dmsoperator.post_compair(revise_file.get("realPath"), edit_res["data"])
        return res["data"]





